// RegSvr.cpp
// (c) 2011, Charles Lechasseur
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#include "stdafx.h"
#include "RegSvr.h"

#include "Exceptions.h"

#include <sstream>


namespace
{
    // Prototype of functions to load from the DLL.
    typedef HRESULT (STDAPICALLTYPE *RegUnregFunctionP)(void);
    typedef HRESULT (STDAPICALLTYPE *InstallFunctionP)(BOOL /*bInstall*/, LPCWSTR /*pszCmdLine*/);

} // anonymous namespace

namespace RG64
{
    //
    // Performs registration actions as appropriate.
    // This function assumes that the DLL is in the correct
    // bitness with regards to the executable platform.
    //
    // @param p_CmdLine Command-line providing options.
    // @return Program return value.
    //
    int PerformRegSvrActions(const CommandLine& p_CmdLine)
    {
        int retValue = 0;

        // Load the DLL into our address space.
        HMODULE hDll = ::LoadLibraryA(p_CmdLine.DllName().c_str());
        if (hDll == 0) {
            throw Win32Exception();
        }

        try {
            // Check if we must register or unregister.
            switch (p_CmdLine.GetRegServerFlag()) {
                case CommandLine::RegServer: {
                    // Call 'DllRegisterServer'.
                    RegUnregFunctionP pFunc = (RegUnregFunctionP) ::GetProcAddress(hDll, "DllRegisterServer");
                    if (pFunc == 0) {
                        throw Win32Exception();
                    }
                    HRESULT hRes = pFunc();
                    retValue = static_cast<int>(hRes);
                    if (!p_CmdLine.IsSilent()) {
                        ATL::CStringA titleStr(MAKEINTRESOURCEA(IDS_APP_TITLE));
                        if (SUCCEEDED(hRes)) {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_REGSERVER_SUCCESS));
                            ::MessageBoxA(0, (LPCSTR) msgStr, (LPCSTR) titleStr, MB_OK | MB_ICONINFORMATION);
                        } else {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_REGSERVER_FAIL));
                            std::string msg((LPCSTR) msgStr);
                            msg += ToHex(hRes);
                            ::MessageBoxA(0, msg.c_str(), (LPCSTR) titleStr, MB_OK | MB_ICONWARNING);
                        }
                    }
                    break;
                }
                case CommandLine::UnRegServer: {
                    // Call 'DllUnregisterServer'.
                    RegUnregFunctionP pFunc = (RegUnregFunctionP) ::GetProcAddress(hDll, "DllUnregisterServer");
                    if (pFunc == 0) {
                        throw Win32Exception();
                    }
                    HRESULT hRes = pFunc();
                    retValue = static_cast<int>(hRes);
                    if (!p_CmdLine.IsSilent()) {
                        ATL::CStringA titleStr(MAKEINTRESOURCEA(IDS_APP_TITLE));
                        if (SUCCEEDED(hRes)) {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_UNREGSERVER_SUCCESS));
                            ::MessageBoxA(0, (LPCSTR) msgStr, (LPCSTR) titleStr, MB_OK | MB_ICONINFORMATION);
                        } else {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_UNREGSERVER_FAIL));
                            std::string msg((LPCSTR) msgStr);
                            msg += ToHex(hRes);
                            ::MessageBoxA(0, msg.c_str(), (LPCSTR) titleStr, MB_OK | MB_ICONWARNING);
                        }
                    }
                    break;
                }
                case CommandLine::NoRegServer:
                default:
                    // Nothing to do here.
                    break;
            }

            // Check if we must install or uninstall.
            switch (p_CmdLine.GetInstallFlag()) {
                case CommandLine::Install: {
                    // Call 'DllInstall' with bInstall = TRUE.
                    InstallFunctionP pFunc = (InstallFunctionP) ::GetProcAddress(hDll, "DllInstall");
                    if (pFunc == 0) {
                        throw Win32Exception();
                    }
                    HRESULT hRes = pFunc(TRUE, p_CmdLine.InstallCmdLine().c_str());
                    retValue = static_cast<int>(hRes);
                    if (!p_CmdLine.IsSilent()) {
                        ATL::CStringA titleStr(MAKEINTRESOURCEA(IDS_APP_TITLE));
                        if (SUCCEEDED(hRes)) {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_INSTALL_SUCCESS));
                            ::MessageBoxA(0, (LPCSTR) msgStr, (LPCSTR) titleStr, MB_OK | MB_ICONINFORMATION);
                        } else {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_INSTALL_FAIL));
                            std::string msg((LPCSTR) msgStr);
                            msg += ToHex(hRes);
                            ::MessageBoxA(0, msg.c_str(), (LPCSTR) titleStr, MB_OK | MB_ICONWARNING);
                        }
                    }
                    break;
                }
                case CommandLine::Uninstall: {
                    // Call 'DllInstall' with bInstall = FALSE.
                    InstallFunctionP pFunc = (InstallFunctionP) ::GetProcAddress(hDll, "DllInstall");
                    if (pFunc == 0) {
                        throw Win32Exception();
                    }
                    HRESULT hRes = pFunc(FALSE, p_CmdLine.InstallCmdLine().c_str());
                    retValue = static_cast<int>(hRes);
                    if (!p_CmdLine.IsSilent()) {
                        ATL::CStringA titleStr(MAKEINTRESOURCEA(IDS_APP_TITLE));
                        if (SUCCEEDED(hRes)) {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_UNINSTALL_SUCCESS));
                            ::MessageBoxA(0, (LPCSTR) msgStr, (LPCSTR) titleStr, MB_OK | MB_ICONINFORMATION);
                        } else {
                            ATL::CStringA msgStr(MAKEINTRESOURCEA(IDS_UNINSTALL_FAIL));
                            std::string msg((LPCSTR) msgStr);
                            msg += ToHex(hRes);
                            ::MessageBoxA(0, msg.c_str(), (LPCSTR) titleStr, MB_OK | MB_ICONWARNING);
                        }
                    }
                    break;
                }
                case CommandLine::NoInstall:
                default:
                    // Nothing to do here.
                    break;
            }
        } catch (...) {
            // Unload DLL before propagating exception.
            ::FreeLibrary(hDll);
            throw;
        }

        // Unload DLL and return.
        ::FreeLibrary(hDll);
        return retValue;
    }

    //
    // Returns a string representation of the given HRESULT as a hex string.
    //
    // @param p_hResult Result to convert
    // @return String with hex representation
    //
    std::string ToHex(const HRESULT p_hResult)
    {
        std::stringstream ss;
        ss << std::hex << p_hResult;
        std::string s = ss.str();
        return s;
    }

} // namespace RG64
